#!/usr/bin/env python3
import sys
import os
import signal
import subprocess
import time
import atexit
import psutil
from pathlib import Path
import threading
import logging

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class ServerManager:
    """服务器管理器"""
    
    def __init__(self):
        self.server_process = None
        self.shutdown_initiated = False
        
        # 注册信号处理器
        signal.signal(signal.SIGTERM, self._signal_handler)
        signal.signal(signal.SIGINT, self._signal_handler)
        
        # 注册退出处理器
        atexit.register(self.cleanup)
    
    def _signal_handler(self, signum, frame):
        """信号处理器"""
        logger.info(f"接收到信号 {signum}，开始关闭...")
        # self.graceful_shutdown()
    
    def check_dependencies(self):
        """检查依赖是否安装"""
        try:
            import fastapi
            import uvicorn
            import sqlalchemy
            import ultralytics
            import psutil
            logger.info("✅ 所有依赖已安装")
            return True
        except ImportError as e:
            logger.error(f"❌ 缺少依赖: {e}")
            print("请运行以下命令安装依赖:")
            print("pip install fastapi uvicorn sqlalchemy ultralytics psutil")
            return False
    
    def setup_environment(self):
        """设置环境"""
        directories = [
            "models/cloud",
            "models/edge", 
            "models/pretrained",
            "logs",
            "temp",
        ]
        
        for dir_path in directories:
            Path(dir_path).mkdir(parents=True, exist_ok=True)
            logger.info(f"✅ 创建目录: {dir_path}")
    
    def start_server(self):
        """启动FastAPI服务器"""
        logger.info("🚀 启动边缘云端AI训练系统...")
        
        # 切换到src目录
        original_cwd = os.getcwd()
        os.chdir("src")
        
        try:
            # 启动服务器
            cmd = [
                sys.executable, "-m", "uvicorn", 
                "main:app", 
                "--host", "0.0.0.0", 
                "--port", "8000",
                "--reload"
            ]
            
            logger.info(f"启动命令: {' '.join(cmd)}")
            self.server_process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                universal_newlines=True,
                bufsize=1
            )
            
            logger.info(f"服务器进程启动，PID: {self.server_process.pid}")
            
            # 启动日志监控线程
            log_thread = threading.Thread(target=self._monitor_server_logs)
            log_thread.daemon = True
            log_thread.start()
            
            # 等待服务器进程
            return_code = self.server_process.wait()
            
            if return_code == 0:
                logger.info("服务器正常退出")
            else:
                logger.error(f"服务器异常退出，返回码: {return_code}")
                
        except KeyboardInterrupt:
            logger.info("接收到中断信号...")
        except Exception as e:
            logger.error(f"启动服务器时出错: {e}")
        finally:
            os.chdir(original_cwd)
    
    def _monitor_server_logs(self):
        """监控服务器日志"""
        if not self.server_process:
            return
            
        try:
            for line in self.server_process.stdout:
                if not self.shutdown_initiated:
                    print(line.strip())
        except Exception as e:
            logger.error(f"监控日志时出错: {e}")
    
    def _is_process_related_to_current_app(self, proc):
        """验证进程是否与当前应用相关"""
        try:
            cmdline = ' '.join(proc.cmdline() or [])
            cmdline_lower = cmdline.lower()
            
            # 严格的匹配条件
            conditions = [
                # 1. 是我们启动的uvicorn服务器
                ('uvicorn' in cmdline_lower and 'main:app' in cmdline_lower),
                
                # 2. 是包含我们特定模块的Python进程
                (sys.executable in cmdline and any(module in cmdline_lower for module in [
                    'ultralytics', 'yolo'
                ]) and os.getcwd() in cmdline),
                
                # 3. 进程的父进程是我们的主进程
                (proc.ppid() == self.server_process.pid if self.server_process else False),
            ]
            
            return any(conditions)
            
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            return False

    def get_all_child_processes(self):
        """获取所有子进程 - 改进版本，更安全"""
        if not self.server_process:
            return []
        
        child_processes = []
        current_user = os.getuid() if hasattr(os, 'getuid') else None
        
        try:
            parent = psutil.Process(self.server_process.pid)
            children = parent.children(recursive=True)
            child_processes.extend(children)
            
            # 只在当前用户权限范围内查找相关进程
            # 使用更精确的匹配条件
            for proc in psutil.process_iter(['pid', 'name', 'cmdline', 'username']):
                try:
                    # 检查进程所有者
                    if current_user is not None:
                        try:
                            proc_uid = proc.uids().real
                            if proc_uid != current_user:
                                continue  # 跳过其他用户的进程
                        except (psutil.NoSuchProcess, psutil.AccessDenied):
                            continue
                    
                    cmdline = ' '.join(proc.info['cmdline'] or [])
                    cmdline_lower = cmdline.lower()
                    
                    # 更精确的匹配条件
                    is_related = (
                        # 检查是否是我们启动的uvicorn服务器
                        ('uvicorn' in cmdline_lower and 'main:app' in cmdline_lower) or
                        # 检查是否是我们的Python进程且包含特定模块
                        (sys.executable in cmdline and any(module in cmdline_lower for module in [
                            'ultralytics', 'yolo'
                        ])) or
                        # 检查是否包含我们的工作目录路径
                        (os.getcwd() in cmdline)
                    )
                    
                    if is_related and proc not in child_processes:
                        # 额外验证：检查进程的工作目录
                        try:
                            proc_cwd = proc.cwd()
                            if proc_cwd and (os.getcwd() in proc_cwd or proc_cwd in os.getcwd()):
                                child_processes.append(proc)
                                logger.info(f"发现相关进程: PID={proc.pid}, 名称={proc.name()}, 命令={cmdline[:100]}...")
                        except (psutil.NoSuchProcess, psutil.AccessDenied):
                            # 如果无法获取工作目录，但其他条件匹配，也添加
                            child_processes.append(proc)
                            logger.info(f"发现相关进程(无法验证工作目录): PID={proc.pid}, 名称={proc.name()}")
                            
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    continue
                    
        except psutil.NoSuchProcess:
            logger.warning("主进程已不存在")
        
        logger.info(f"总共发现 {len(child_processes)} 个相关进程")
        return child_processes
    
    def kill_process_tree(self, timeout=30):
        """终止进程树"""
        if not self.server_process:
            return
        
        logger.info("正在终止服务器进程树...")
        
        try:
            # 获取所有子进程
            child_processes = self.get_all_child_processes()
            
            # 首先尝试关闭主进程
            if self.server_process.poll() is None:
                logger.info(f"发送SIGTERM到主进程 {self.server_process.pid}")
                self.server_process.terminate()
            
            # 等待主进程退出
            try:
                self.server_process.wait(timeout=timeout//2)
                logger.info("主进程已退出")
            except subprocess.TimeoutExpired:
                logger.warning("主进程超时，强制终止")
                self.server_process.kill()
                self.server_process.wait(timeout=5)
            
            # 处理子进程
            if child_processes:
                logger.info(f"发现 {len(child_processes)} 个子进程，正在终止...")
                
                # 首先尝试SIGTERM
                for proc in child_processes:
                    try:
                        logger.info(f"终止子进程: PID={proc.pid}, 名称={proc.name()}")
                        proc.terminate()
                    except (psutil.NoSuchProcess, psutil.AccessDenied):
                        continue
                
                # 等待子进程退出
                start_time = time.time()
                remaining_processes = child_processes[:]
                
                while remaining_processes and (time.time() - start_time) < timeout//2:
                    for proc in remaining_processes[:]:
                        try:
                            if proc.poll() is not None or not proc.is_running():
                                remaining_processes.remove(proc)
                        except (psutil.NoSuchProcess, psutil.AccessDenied):
                            remaining_processes.remove(proc)
                    time.sleep(1)
                
                # 强制终止剩余进程
                if remaining_processes:
                    logger.warning(f"强制终止 {len(remaining_processes)} 个剩余进程")
                    for proc in remaining_processes:
                        try:
                            proc.kill()
                            proc.wait(timeout=3)
                        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.TimeoutExpired):
                            continue
            
            logger.info("进程树终止完成")
            
        except Exception as e:
            logger.error(f"终止进程树时出错: {e}")
    
    def graceful_shutdown(self):
        """关闭"""
        if self.shutdown_initiated:
            return
        
        self.shutdown_initiated = True
        logger.info("开始关闭...")
        
        try:
            # 发送关闭请求到API（如果服务器还在运行）
            try:
                import requests
                response = requests.post("http://localhost:8000/api/system/shutdown", timeout=5)
                logger.info("已发送关闭请求到API")
                time.sleep(2)  # 给API一些时间处理
            except:
                logger.info("无法连接到API，直接终止进程")
            
            # 终止进程树
            # self.kill_process_tree(timeout=30)
            
            # 清理临时文件
            # self.cleanup_temp_files()
            
            logger.info("关闭完成")
            
        except Exception as e:
            logger.error(f"关闭时出错: {e}")
    
    def cleanup_temp_files(self):
        """清理临时文件"""
        try:
            temp_dir = Path("temp")
            if temp_dir.exists():
                import shutil
                for item in temp_dir.iterdir():
                    try:
                        if item.is_dir():
                            shutil.rmtree(item)
                        else:
                            item.unlink()
                        logger.info(f"清理临时文件: {item}")
                    except Exception as e:
                        logger.warning(f"清理文件失败 {item}: {e}")
        except Exception as e:
            logger.error(f"清理临时文件时出错: {e}")
    
    def cleanup(self):
        """最终清理"""
        # 暂时不关闭
        ...
        # if not self.shutdown_initiated:
            # self.graceful_shutdown()

def main():
    """主函数"""
    
    server_manager = ServerManager()
    
    if not server_manager.check_dependencies():
        sys.exit(1)
    
    server_manager.setup_environment()
    
    try:
        server_manager.start_server()
    except KeyboardInterrupt:
        logger.info("接收到键盘中断")
    finally:
        logger.info("关闭完成")
        # server_manager.graceful_shutdown()

if __name__ == "__main__":
    main()